import os
import sys
import onnx
import onnx_graphsurgeon as gs
import numpy as np

PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

ONNX_BEFORE = os.path.join(PROJECT_ROOT, "onnx/nms_before_q.onnx")
ONNX_AFTER  = os.path.join(PROJECT_ROOT, "onnx/nms_after.onnx")
ONNX_MERGED = os.path.join(PROJECT_ROOT, "onnx/nms_merged_q.onnx")

before_graph = gs.import_onnx(onnx.load(ONNX_BEFORE))
after_graph  = gs.import_onnx(onnx.load(ONNX_AFTER))

# 1. 获取backbone输出和NMS输入
backbone_output_var = before_graph.outputs[0]
nms_input_var = after_graph.inputs[0]

print(f"Backbone输出: {backbone_output_var.name}")
print(f"NMS输入: {nms_input_var.name}")

# 2. 直接替换NMS子图的输入变量
# 将after_graph的输入替换为backbone的输出
after_graph.inputs[0] = backbone_output_var

# 3. 更新所有引用原NMS输入变量的节点
for node in after_graph.nodes:
    # 替换节点的输入引用
    for i, inp in enumerate(node.inputs):
        if inp is nms_input_var or inp.name == nms_input_var.name:
            node.inputs[i] = backbone_output_var

# 4. 设置最终的双输出：NMS结果 + backbone特征
merged_outputs = list(after_graph.outputs)  # NMS的输出
merged_outputs.append(backbone_output_var)   # backbone特征输出

# 5. 合并所有组件（不需要Identity节点了）
merged_nodes = list(before_graph.nodes) + list(after_graph.nodes)
merged_inputs = before_graph.inputs

# 6. 合并常量
merged_tensors = {}
for t in before_graph.tensors().values():
    if isinstance(t, gs.Constant):
        merged_tensors[t.name] = t
for t in after_graph.tensors().values():
    if isinstance(t, gs.Constant):
        merged_tensors[t.name] = t

merged_graph = gs.Graph(
    nodes=merged_nodes,
    inputs=merged_inputs,
    outputs=merged_outputs
)
for t in merged_tensors.values():
    merged_graph.tensors()[t.name] = t

# 7. 修复TensorRT不兼容的节点
def fix_unsqueeze_nodes(graph):
    tensor_map = {t.name: t for t in graph.tensors().values()}
    for node in graph.nodes:
        if node.op == "Unsqueeze" and len(node.inputs) == 2:
            data_input = node.inputs[0]
            axes_input = node.inputs[1]
            axes_value = None
            if isinstance(axes_input, gs.Constant):
                axes_value = axes_input.values
            elif axes_input.name in tensor_map and hasattr(tensor_map[axes_input.name], "values"):
                axes_value = tensor_map[axes_input.name].values
            if axes_value is not None:
                node.attrs['axes'] = axes_value.tolist() if hasattr(axes_value, 'tolist') else list(axes_value)
                node.inputs = [data_input]

def fix_squeeze_nodes(graph):
    tensor_map = {t.name: t for t in graph.tensors().values()}
    for node in graph.nodes:
        if node.op == "Squeeze" and len(node.inputs) == 2:
            data_input = node.inputs[0]
            axes_input = node.inputs[1]
            axes_value = None
            if isinstance(axes_input, gs.Constant):
                axes_value = axes_input.values
            elif axes_input.name in tensor_map and hasattr(tensor_map[axes_input.name], "values"):
                axes_value = tensor_map[axes_input.name].values
            if axes_value is not None:
                node.attrs['axes'] = axes_value.tolist() if hasattr(axes_value, 'tolist') else list(axes_value)
                node.inputs = [data_input]

def fix_split_nodes(graph):
    tensor_map = {t.name: t for t in graph.tensors().values()}
    for node in graph.nodes:
        if node.op == "Split" and len(node.inputs) == 2:
            data_input = node.inputs[0]
            split_input = node.inputs[1]
            split_value = None
            if isinstance(split_input, gs.Constant):
                split_value = split_input.values
            elif split_input.name in tensor_map and hasattr(tensor_map[split_input.name], "values"):
                split_value = tensor_map[split_input.name].values
            if split_value is not None:
                node.attrs['split'] = split_value.tolist() if hasattr(split_value, 'tolist') else list(split_value)
                node.inputs = [data_input]

fix_unsqueeze_nodes(merged_graph)
fix_squeeze_nodes(merged_graph)
fix_split_nodes(merged_graph)

# 8. 清理和保存
merged_graph.cleanup().toposort()
onnx.save(gs.export_onnx(merged_graph), ONNX_MERGED)

print("合成成功！无Identity节点的直连方式")
print(f"输出0 (NMS结果): {merged_outputs[0].name}")
print(f"输出1 (Backbone特征): {merged_outputs[1].name}")
print("数据流: Backbone → 直接连接 → NMS")